import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from torch.distributions import Categorical
import time
import gym

"""Pytorch中神经网络模块化接口nn的了解"""
"""
torch.nn是专门为神经网络设计的模块化接口。nn构建于autograd之上，可以用来定义和运行神经网络。
nn.Module是nn中十分重要的类,包含网络各层的定义及forward方法。
定义自已的网络：
    需要继承nn.Module类，并实现forward方法。
    一般把网络中具有可学习参数的层放在构造函数__init__()中，
    不具有可学习参数的层(如ReLU)可放在构造函数中，也可不放在构造函数中(而在forward中使用nn.functional来代替)

    只要在nn.Module的子类中定义了forward函数，backward函数就会被自动实现(利用Autograd)。
    在forward函数中可以使用任何Variable支持的函数，毕竟在整个pytorch构建的图中，是Variable在流动。还可以使用
    if,for,print,log等python语法.

    注：Pytorch基于nn.Module构建的模型中，只支持mini-batch的Variable输入方式，
    比如，只有一张输入图片，也需要变成 N x C x H x W 的形式：

    input_image = torch.FloatTensor(1, 28, 28)
    input_image = Variable(input_image)
    input_image = input_image.unsqueeze(0)   # 1 x 1 x 28 x 28

"""
class PolicyGradientNetwork(nn.Module): # 定义了actor的结构
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(8, 16) # fc是全连接层。这里输入是8维的，是env输出的observation，输出16维
        self.fc2 = nn.Linear(16, 16) # 输入是16维的，输出16维
        self.fc3 = nn.Linear(16, 4) # 输入是16维的，输出4维，代表了四种行为

    def forward(self, state):
        hid = torch.tanh(self.fc1(state)) # tanh是激活函数
        hid = torch.tanh(self.fc2(hid))
        return F.softmax(self.fc3(hid), dim=-1) # softmax函数把值转化为概率


class PolicyGradientAgent():
    def __init__(self, network):
        self.network = network
        # optimizer用怎样的策略更新参数，主要为了加快训练速度，减少过拟合
        self.optimizer = optim.SGD(self.network.parameters(), lr=0.000001)

    def learn(self, log_probs, rewards):
        loss = (-log_probs * rewards).sum() # 每一步的概率 * 最终的reward，然后相加，这是policy gradient的做法

        self.optimizer.zero_grad()
        loss.backward()
        self.optimizer.step()

    def sample(self, state):
        action_prob = self.network(torch.FloatTensor(state)) # 这一步把数据放到网络中，得到结果，也即四种行为的概率
        # 作用是创建以参数probs为标准的类别分布，样本是来自 “0 … K - 1” 的整数，其中
        # K 是probs参数的长度。也就是说，按照传入的probs中给定的概率，在相应的位置处进行取样，取样返回的是该位置的整数索引。
        # 如果 probs 是长度为 K 的一维列表，则每个元素是对该索引处的类进行抽样的相对概率。
        # 如果 probs 是二维的，它被视为一批概率向量。
        action_dist = Categorical(action_prob)
        action = action_dist.sample()
        log_prob = action_dist.log_prob(action)
        return action.item(), log_prob

if __name__ == '__main__':
    env = gym.make('LunarLander-v2')

    # 测试随机是否能成功
    env.reset()
    done = False
    while not done:
        action = env.action_space.sample() # 这个sample返回的是数字
        observation, reward, done, _ = env.step(action)
    print(reward)

    network = PolicyGradientNetwork()
    agent = PolicyGradientAgent(network)

    agent.network.load_state_dict(torch.load('rl_round.pkl'))

    agent.network.train()  # 訓練前，先確保 network 處在 training 模式
    EPISODE_PER_BATCH = 5  # 每蒐集 5 個 episodes 更新一次 agent
    NUM_BATCH = 40000  # 總共更新 400 次
    avg_total_rewards, avg_final_rewards = [], []

    for batch in range(NUM_BATCH):
        log_probs, rewards = [], []
        total_rewards, final_rewards = [], []

        epoch_start_time = time.time()

        # 蒐集訓練資料
        for episode in range(EPISODE_PER_BATCH):

            state = env.reset()
            total_reward, total_step = 0, 0

            while True:
                action, log_prob = agent.sample(state) # 返回action和概率的log值
                next_state, reward, done, _ = env.step(action)

                log_probs.append(log_prob)
                state = next_state
                total_reward += reward
                total_step += 1

                if done:
                    final_rewards.append(reward)
                    total_rewards.append(total_reward)
                    rewards.append(np.full(total_step, total_reward - 160))  # 設定同一個 episode 每個 action 的 reward 都是 total reward
                    break

        # 紀錄訓練過程
        avg_total_reward = sum(total_rewards) / len(total_rewards)
        avg_final_reward = sum(final_rewards) / len(final_rewards)
        avg_total_rewards.append(avg_total_reward)
        avg_final_rewards.append(avg_final_reward)
        print("Batch {},\t{:.2f} sec(s)\tTotal Reward = {:.1f},\tFinal Reward = {:.1f} \tstep {}".format(
            batch + 1, time.time() - epoch_start_time, avg_total_reward, avg_final_reward, total_step))

        # 更新網路
        rewards = np.concatenate(rewards, axis=0)
        rewards = (rewards - np.mean(rewards)) / (np.std(rewards) + 1e-9)  # 將 reward 正規標準化
        agent.learn(torch.stack(log_probs), torch.from_numpy(rewards))
        torch.save(agent.network.state_dict(), 'rl_round.pkl')